import torch
import numpy as np
from tqdm import tqdm
from DDPM.Config import *
from DDPM.Evaluation.Control import *
from DDPM.Evaluation.InternalBallistic import get_target_from_generated_samples ,ballistic_curve_from_index, get_av2_values_from_generated_samples,convert_av2_to_value
from Utility.PlotCurve import plot_curves
from DDPM.Evaluation.Decoding import decode_all
import shutil

def load_data_psfd(control_params, data_path: str, data_file_num:int, data_index_start:int=0):  # 载入PSFD数据
    from OPT.Objective import performance_matching_degree
    project_name=control_params["project_name"] 
    user_defined=control_params["user_defined"]
    
    project_path=f"{decode_result_save_path}/{project_name}"
    if user_defined:
        t_array_target, p_array_target = get_target_from_generated_samples(project_path)
        cstar, rho, n, At = get_av2_values_from_generated_samples(project_path)
        cstar_rho = cstar * rho
    else:
        i_from_psfd=control_params["i_from_psfd"]
        t_array_target, p_array_target = ballistic_curve_from_index(trainingset_path, i_from_psfd, 0)


    t_array_target_tensor = torch.tensor(t_array_target, dtype=torch.float32)
    t_max=t_array_target_tensor.max()
    p_array_target_tensor = torch.tensor(p_array_target, dtype=torch.float32)

    results = []  # 用于存储 (i, j, r2) 的列表
    for i in tqdm(range(data_index_start, data_index_start + data_file_num), desc='loading data', colour='green', dynamic_ncols=True):
        npz_path = f'{data_path}/{i}'
        npz_file = np.load(npz_path+".npz")
        bvs = npz_file["bvs"]
        avs2 = npz_file["avs2"]
        item_num_per_file = npz_file["avs2"].shape[0]

        for j in range(item_num_per_file):
            t_array_designed_tensor = torch.tensor(bvs[j,:,0], dtype=torch.float32)
            p_array_designed_tensor = torch.tensor(bvs[j,:,1], dtype=torch.float32)*1e6
            
            cstar_rho_psfd, n_psfd, At_psfd=convert_av2_to_value(avs2[j])

            #只在PSFD中查询满足cstar_rho，n，At约束的解（误差正负5%）
            w=0.1
            if (1-w)*cstar_rho< cstar_rho_psfd and cstar_rho_psfd<(1+w)*cstar_rho\
                and (1-w)*n < n_psfd and n_psfd<(1+w)*n\
                and (1-w)*At< At_psfd and At_psfd<(1+w)*At\
                and (1-w)*t_max< t_array_designed_tensor.max() and t_array_designed_tensor.max()<(1+w)*t_max:
                with torch.no_grad():
                    r2, r2_k, t_vector_designed, p_vector_designed=\
                        performance_matching_degree(t_array_designed_tensor,
                                                p_array_designed_tensor,
                                                t_array_target_tensor,
                                                p_array_target_tensor)
                    results.append((i, j, r2, r2_k))  # 将 (i, j, r2) 存入列表
    # 按照 r2 从大到小排序
    results.sort(key=lambda x: x[2], reverse=True)

    # 输出排序后的前5个解
    for result in results[:5]:
        print(f"i={result[0]}, j={result[1]}, r2={result[2]},r2_k={result[3]}")

    os.makedirs(f"{project_path}/PSFD", exist_ok=True)
    np.savez_compressed(f"{project_path}/PSFD/internal_ballistic_psfd.npz",
                        i_list=[r[0] for r in results],
                        j_list=[r[1] for r in results],
                        r2_list=[r[2] for r in results],
                        r2_k_list=[r[3] for r in results])

    return results

def start_find_best_in_psfd(control_params):
    project_name=control_params["project_name"]
    t_array_target=control_params["t_array_target"]
    p_array_target=control_params["p_array_target"]

    project_path=f"{decode_result_save_path}/{project_name}"

    # 清空目录
    if os.path.exists(f"{project_path}/PSFD"):
        shutil.rmtree(f"{project_path}/PSFD")
    os.makedirs(f"{project_path}/PSFD", exist_ok=True)

    results=load_data_psfd(control_params, trainingset_path, data_file_num)

    t_array_list=[]
    p_array_list=[]

    lvs=[]
    gvs=[]

    for result in results[:5]:
        i=result[0]
        j=result[1]
        npz_path = f'{trainingset_path}/{i}'
        npz_file = np.load(npz_path+".npz")
        bvs = npz_file["bvs"]
        lv= npz_file["lv"]
        gv= npz_file["gv"]
        lvs.append(lv)
        gvs.append(gv)

        item_num_per_file = npz_file["avs2"].shape[0]

        t_array_list.append(bvs[j,:,0])
        p_array_list.append(bvs[j,:,1]*1e6)
    

    # 保存PSFD中查询到的样本到新的 npz 文件
    save_path = f"{project_path}/PSFD/Samples_best.npz"
    np.savez_compressed(save_path,
                        lv=np.array(lvs),
                        gv=np.array(gvs))
    
    decode_all(f"{project_path}/PSFD","Samples_best")

    # 绘制曲线图
    plot_curves(title='Chamber Pressure vs Burned time',
        xlabel='$\it{t}$ /s',  # t 设置为斜体，保留字体
        ylabel='$\it{p}$ /MPa',  # p 设置为斜体，保留字体
        x=np.array(t_array_list).T,
        y=np.array(p_array_list).T * 1e-6,
        ref_x=t_array_target,
        ref_y=p_array_target * 1e-6,
        filename=f"{project_path}/PSFD/internal_ballistic_psfd.svg")

if __name__ == '__main__':
    control_params=get_control_params(question_name="试验2",
        postfix="_epoch50000",
        cluster_num=5,
        guid_w=0.5)
    start_find_best_in_psfd(control_params)